package com.leetcode.程序员面试金典.hard;

import java.util.ArrayList;
import java.util.Collections;
import java.util.function.BiPredicate;

/**
 * @author 覃国强
 * @date 2022-02-16 15:45
 */
public class 面试题_17_20_连续中值 {

  public static void main(String[] args) {
    MedianFinder obj = new MedianFinder();
    obj.addNum(1);
    obj.addNum(2);
    System.out.println(obj.findMedian());
    obj.addNum(3);
    System.out.println(obj.findMedian());
  }

  static
  class MedianFinder {

    private final ArrayList<Integer> largeTopHeap = new ArrayList<>(Collections.singletonList(Integer.MIN_VALUE));
    private final BiPredicate<Integer, Integer> largeTopHeapPredicate = (cNum, fNum) -> cNum > fNum;


    private final ArrayList<Integer> smallTopHeap = new ArrayList<>(Collections.singletonList(Integer.MAX_VALUE));
    private final BiPredicate<Integer, Integer> smallTopHeapPredicate = (cNum, fNum) -> cNum < fNum;

    /** initialize your data structure here. */
    public MedianFinder() {

    }

    public void addNum(int num) {
      // 小于等于右区间最小值，加入左区间
      if (num <= smallTopHeap.get(0)) {
        addHeap(num, largeTopHeap, largeTopHeapPredicate);
      } else {
        addHeap(num, smallTopHeap, smallTopHeapPredicate);
      }

      // 左区间 - 右区间 > 1
      if (largeTopHeap.size() - smallTopHeap.size() > 1) {
        // 删除左区间
        Integer top = removeHeap(largeTopHeap, largeTopHeapPredicate);
        // 加入右区间
        addHeap(top, smallTopHeap, smallTopHeapPredicate);
      }

      // 右区间 - 左区间 > 0
      if (smallTopHeap.size() - largeTopHeap.size() > 1) {
        // 删除右区间
        Integer top = removeHeap(smallTopHeap, smallTopHeapPredicate);
        // 加入左区间
        addHeap(top, largeTopHeap, largeTopHeapPredicate);
      }
    }

    public double findMedian() {
      if (smallTopHeap.size() > largeTopHeap.size()) {
        return smallTopHeap.get(0);
      } else if (largeTopHeap.size() > smallTopHeap.size()) {
        return largeTopHeap.get(0);
      } else {
        return (largeTopHeap.get(0) + smallTopHeap.get(0)) / 2.0;
      }
    }

    private void addHeap(int num, ArrayList<Integer> heap, BiPredicate<Integer, Integer> predicate) {
      heap.add(num);
      adjustHeap(heap, predicate);
    }

    private Integer removeHeap(ArrayList<Integer> heap, BiPredicate<Integer, Integer> predicate) {
      int top = heap.get(0);
      int end = heap.size() - 1;
      heap.set(0, heap.remove(end));
      int f = 0;
      while (f < heap.size() - 1) {
        int c1 = 2 * f + 1;
        if (c1 > heap.size() - 1) {
          break;
        }
        int c2 = 2 * f + 2;
        if (c2 > heap.size() - 1) {
          c2 = c1;
        }

        int c = getTargetIndex(heap, c1, c2, predicate);
        if (predicate.test(heap.get(c), heap.get(f))) {
          swap(heap, f, c);
        } else {
          break;
        }
        f = c;
      }
      return top;
    }

    private void adjustHeap(ArrayList<Integer> heap, BiPredicate<Integer, Integer> predicate) {
      int c = heap.size() - 1;
      while (c > 0) {
        int f = (c - 1) / 2;
        if(predicate.test(heap.get(c), heap.get(f))) {
          swap(heap, c, f);
        } else {
          break;
        }
        c = f;
      }
    }

    private void swap(ArrayList<Integer> heap, int i, int j) {
      int tmp = heap.get(i);
      heap.set(i, heap.get(j));
      heap.set(j, tmp);
    }

    private int getTargetIndex(ArrayList<Integer> heap, int c1, int c2, BiPredicate<Integer, Integer> predicate) {
      return predicate.test(heap.get(c1), heap.get(c2)) ? c1 : c2;
    }
  }

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

}
